package com.hy.day;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * Created With IntelliJ IDEA.
 * Descriptions:3067. 在带权树网络中统计可连接服务器对数目
 * <p>
 * User:Mr.Du
 * Date:2024/6/4
 * Time:9:26
 */
public class CountPairsOfConnectableServers {

    /**
     * 计算每台服务器可以连接的其他服务器对数。
     * 通过构建服务器之间的图关系，并使用深度优先搜索（DFS）来确定连接可能性。
     *
     * @param edges 服务器之间的连接边，每个元素是一个数组，包含两个服务器编号和连接权重。
     * @param signalSpeed 信号传输速度，用于确定服务器之间是否可以连接。
     * @return 每台服务器可以连接的其他服务器对数的数组。
     */
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        // 服务器数量比边的数量多1，因为边是服务器之间的关系。
        int n = edges.length + 1;
        // 使用邻接表表示图的关系，g[i]表示服务器i可以连接到的服务器列表。
        List<int[]>[] g = new ArrayList[n];
        // 初始化邻接表，确保每个服务器都有一个空的连接列表。
        Arrays.setAll(g, i -> new ArrayList<>());
        // 遍历边，构建无向图的邻接表。
        for (int[] e : edges) {
            int x = e[0];
            int y = e[1];
            int wt = e[2];
            // 由于是无向图，每条边的两个服务器都要互相记录。
            g[x].add(new int[]{y, wt});
            g[y].add(new int[]{x, wt});
        }
        // ans数组用于存储每个服务器可以连接的其他服务器对数。
        int[] ans = new int[n];
        // 遍历每个服务器，计算它可以连接的服务器对数。
        for (int i = 0; i < n; i++) {
            int sum = 0;
            // 遍历服务器i可以连接的其他服务器。
            for (int[] e : g[i]) {
                int cnt = dfs(e[0], i, e[1], g, signalSpeed);
                // 累加到当前服务器的连接对数。
                ans[i] += cnt * sum;
                // 更新已经计算出的连接服务器的数量。
                sum += cnt;
            }
        }
        // 返回每个服务器可以连接的服务器对数。
        return ans;
    }

    /**
     * 深度优先搜索函数，用于计算满足特定条件的子树数量。
     *
     * @param x 当前节点。
     * @param fa 父节点。
     * @param sum 当前路径上的信号强度总和。
     * @param g 图的邻接表表示。
     * @param signalSpeed 信号速度。
     * @return 返回以当前节点为根的子树中，满足信号强度总和能被信号速度整除的子树数量。
     */
    private int dfs(int x, int fa, int sum, List<int[]>[] g, int signalSpeed) {
        // 判断当前路径上的信号强度总和是否能被信号速度整除，如果可以，则计数加1。
        int cnt = sum % signalSpeed == 0 ? 1 : 0;
        // 遍历当前节点的所有子节点。
        for (int[] e : g[x]) {
            int y = e[0];
            // 排除父节点，避免重复计算。
            if (y != fa) {
                // 递归计算子节点为根的子树中满足条件的子树数量，并累加到计数中。
                cnt += dfs(y, x, sum + e[1], g, signalSpeed);
            }
        }
        // 返回满足条件的子树数量。
        return cnt;
    }

    public static void main(String[] args) {
        int[][] edges = {
                {0, 6, 3},
                {6, 5, 3},
                {0, 3, 1},
                {3, 2, 7},
                {3, 1, 6},
                {3, 4, 2}
        };
        CountPairsOfConnectableServers solution = new CountPairsOfConnectableServers();
        int[] res = solution.countPairsOfConnectableServers(edges, 3);
        System.out.println(Arrays.toString(res));
    }

}
